"""
File         lu.py
Author    Ernesto P. Adorio
               (UPEPP) U.P. at Clarkfield 
               Pampanga, the Philippines
email:      ernesto. adorio @gmail .com
Description
             LU routines in Python
 
References
      Conte and deBoor "Elementary Numerical Analysis", 3e, pp.Sec. 4.,  pp.160-168.
 
Revisions
    2009.01.22    First Release.
"""
 
#from matlib import *
 
def matlufactor(A, ztol = 1.0e-12):
     """
      ipivot, LU = matlufac(A)
 
     Input arguments
          A                 input: matrix to be factored
          ztol              input: zero value tolerance
 
     Output
          ipivot           output:  array of row interchanges made.
          LU                output:  LU decomposition of matrix A.
         detsign         output:  Sign of determinant.          
 
     Description
          Splits a square matrix A into  a lower and an upper
          triangular matrix L and U, stored in LU, with the row interchanges
          stored in  ipivot using implicit scaled partial pivoting.
 
     Return value
          Sign of determinant.
     Reference
          Based on a FORTRAN 77 routine of Conte and de Boor.
     Revisions
        2009.01.23       debugged and released.
     """
     (nrows,  ncols)   = [len(A),len(A[0])]
     LU        = [A[i][:] for i in range(nrows)]
     ipivot    = range(nrows)
     D          = [1.0]* nrows
     detsign = 1
 
     #Fatoring matrix.
     # initialize rowperm, D, detsign
     D = [0.0] * nrows
     detsign = 1
     for i in range(nrows):
          ipivot[i] = i
          rowmax = max([abs(x) for x in LU[i]])
          if rowmax <= ztol: 
             detsign =  0
             rowmax  = 1.0
          D[i] = rowmax
     if nrows <= 1:
          return None,  None,  None
 
     for k in range(nrows-1):
          colmax   = abs(LU[k][k]) / D[k]
          istar = k
          for i in range(k+1, nrows):
              t = abs(LU[i][k]) / D[i]
              if t > colmax:
                 colmax  = t
                 istar = i
          if colmax <= ztol:
              detsign = 0
 
          else:
             if istar > k:
                 #Make row exchanges?
                 detsign                    = -detsign
                 ipivot[istar], ipivot[k] = ipivot[k], ipivot[istar]
                 D[istar], D[k] = D[k], D[istar]
                 for j in range(ncols):
                    LU[istar][j], LU[k][j] = LU[k][j],LU[istar][j]
 
             # Elimination
             for i in range(k+1, nrows):
                  ratio = LU[i][k] = LU[i][k] / LU[k][k]
                  for j in range(k+1,  ncols):
                     LU[i][j] = LU[i][j] - ratio * LU[k][j]
 
     return LU, ipivot, detsign
 
 
def matlusolve(LU, ipivot,B):
      """
      Parameters
          LU           LU decomposition matrix
          ipivot      row interchanges vector
          B             RHS vector of (LU) X = B
      Output:
          X             unknown solution vector to be solved
      Description
          Solves the equation  (LU) X = B for X.
          Once we have an LU factorization of a matrix A, 
          it is very easy to solve for X given any RHS vector B.
      Reference
          This routine is based on a FORTRAN 77 routine
          SUBST of Conte and deBoor.
      """
      nrows,  ncols = [len(LU),len(LU[0])]
      X                    = B[:]
      if (nrows == 1) :
            X[0] = B[0] / LU[0][0]
            return X
 
      #  Forward substitution
      X[0] = B[ipivot[0]]
      for i in range(1, nrows):
         t = 0.0
         for j in range(i):
            t+= (LU[i][j] * X[j])
         X[i] = B[ipivot[i]] - t
 
      # Back substitution
      X[nrows-1] = X[nrows-1] / LU[nrows-1][nrows-1]
      for i in range(nrows -2 , -1, -1):
          t = sum([LU[i][j] * X[j] for j in range(i+1,  ncols)])
          X[i] = (X[i] - t) / LU[i][i]
      return X
 
def  matluinv(LU, ipivot):
    """
    Parameters
          LU                input:  LU decomposition of a matrix A
          InvA             output: returned inverse of A
          ipivot           input:  row interchanges vector
 
    Return value
          0                 normal success
 
    Description
          Computes the inverse of matrix LU in InvA. Because the LU decompositon
          is performed before calling this routine, matrix A or InvA should be
          non-singular.
 
    Description
          Computes InvA = inv(A) by using the LU decomposition of  A.
          LU and ipivot are set by a call to matLUfac().
    """
    nrows = len(LU)
    X    = [0.0] * nrows
    Ainv = []
    #make repeated calls to matlusolve.
    for i in range(nrows):
         # Set up identity matrix column
         B   = [0.0] * nrows
         B[i] = 1.0
 
         X = matlusolve(LU, ipivot, B)
         Ainv.append(X[:])
    return zip(*Ainv)
 
def matinv(A):
    """
    Convenience routine for matrix inverses using input matrix A only as the argument.
    """
    LU, ipivot, detsign  = matlufactor(A, ztol = 1.0e-12)
    return matluinv(LU, ipivot)
 
def matiterimprov(LU, rowPerm, A, Xo, b, maxiter = 1):
     """
     Not thoroughly tested!
 
     LU           - deccomposition of matrix A.
     rowPerm - row interchanges integer vector
     Xo           -  initial estimate of solution vector
     b             - known righthand side vector of AX = b
     maxtier   - maximum number of iterations for improvement.
 
     Pseudocode:
 
     do until satisfied:
       r = b - Ax
       Solve for e in LU e = r
       X = X + e
    Reference: Conte and deBoor, Numerical Analysis.
     """
     nrows = len(A)
     Xnew = Xo[:]
     for i in range(maxiter+1):
         Xnew = matvec(A, Xnew)
         r = [b[i] - Xnew[i] for i in range(nrows)]
         e = matlusolve(LU,rowPerm, b)
         Xnew = [Xnew[i] + e[i] for i in range(nrow)]
     return Xnew
 
 
 
def Test():
    A = [[3.0, 5.0, 2.0], 
             [0.0, 8.0, 2.0], 
             [6.0, 2.0, 8.0]]
    print "Input matrix A:"
    print (A) 
    LU,  ipivot,  detsign = matlufactor(A)
    print "Output LU factorization:"
    print(LU) 
    print "Output Row pivot vector:",  ipivot
    print "Output Sign of determinant:",  detsign
    invLU = matluinv(LU,  ipivot)
    print "matrix inverse using matluinv."
    print(invLU)
    B = matmul(invLU,  A)
    print "(invLU A) Should be the identity."
    print (B) 
    B = matmul(A,  invLU)
    print "(A invLU) Should be the identity."
    print (B) 
    print
 
if __name__ == "__main__":
    Test()
